package cn.seaboot.commons.thread;

import cn.seaboot.commons.exception.HandlerException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.Closeable;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.BiConsumer;
import java.util.function.Consumer;
import java.util.function.Supplier;

/**
 * 一个简单的工作任务
 * <p>
 * 将一个线程产生的数据，随机分配到多个子线程。
 *
 * @param <T> 上下文-数据处理器
 * @param <U> 待处理数据
 * @author Mr.css
 * @version 2025-05-07 14:13
 */
public class WorkTask<T, U> implements Closeable {
    private final Logger logger = LoggerFactory.getLogger(WorkTask.class);

    /**
     * 计时器
     * <p>
     * 当一个工作线程结束任务之后计数减一
     */
    private AtomicInteger latch;

    /**
     * 启用多少个处理线程
     */
    private int size;
    /**
     * 消息队列
     */
    private BlockingQueue<Object> queue;

    /**
     * 上下文生成函数
     * <p>
     * 数据处理函数，可以是一个数据库连接
     */
    private Supplier<T> supplier;
    /**
     * 上下文释放函数
     */
    private Consumer<T> release;

    /**
     * 生产者
     * <p>
     * 一般是是个数据库查询
     */
    private Consumer<WorkTask<T, U>> provider;

    /**
     * 消费者
     */
    private BiConsumer<T, U> consumer;

    /**
     * 子线程异常处理程序
     * <p>
     * 一般用于记录日志，如果继续抛出异常，则会在抛出异常之后，中断工作线程
     */
    private BiConsumer<Exception, U> fullback;
    /**
     * 任务完成回调
     * <p>
     * 所有工作线程完成之后，会自动调用当前函数，一般用于资源回收
     */
    private Consumer<WorkTask<T, U>> finalized;

    /**
     * 执行程序
     */
    private final List<WorkRunnable<T, U>> list = new ArrayList<>();

    public int getSize() {
        return size;
    }

    public void setSize(int size) {
        this.size = size;
    }

    public BlockingQueue<Object> getQueue() {
        return queue;
    }

    public void setQueue(BlockingQueue<Object> queue) {
        this.queue = queue;
    }

    public Supplier<T> getSupplier() {
        return supplier;
    }

    public void setSupplier(Supplier<T> supplier) {
        this.supplier = supplier;
    }

    public Consumer<T> getRelease() {
        return release;
    }

    public void setRelease(Consumer<T> release) {
        this.release = release;
    }

    public BiConsumer<T, U> getConsumer() {
        return consumer;
    }

    public void setConsumer(BiConsumer<T, U> consumer) {
        this.consumer = consumer;
    }

    public Consumer<WorkTask<T, U>> getProvider() {
        return provider;
    }

    public void setProvider(Consumer<WorkTask<T, U>> provider) {
        this.provider = provider;
    }

    public BiConsumer<Exception, U> getFullback() {
        return fullback;
    }

    public void setFullback(BiConsumer<Exception, U> fullback) {
        this.fullback = fullback;
    }

    public Consumer<WorkTask<T, U>> getFinalized() {
        return finalized;
    }

    public void setFinalized(Consumer<WorkTask<T, U>> finalized) {
        this.finalized = finalized;
    }

    /**
     * 初始函数
     */
    private void init() {
        // default exception handler
        if (this.fullback == null) {
            this.fullback = (e, u) -> {
                logger.error("{} processing failed: {}", Thread.currentThread(), u);
            };
        }
    }

    /**
     * create runnable
     *
     * @return working runnable
     */
    private WorkRunnable<T, U> runnable() {
        T ctx = supplier == null ? null : supplier.get();

        WorkRunnable<T, U> run = new WorkRunnable<>(this);
        run.setContext(ctx);
        run.setQueue(queue);
        run.setRelease(release);
        run.setConsumer(consumer);
        run.setFullback(fullback);
        this.list.add(run);
        return run;
    }

    /**
     * 查询数据并加入到队列中
     * <p>
     * 由 current-thread 执行
     */
    private void provide() {
        this.provider.accept(this);
        this.putDumb();
    }

    /**
     * 通过 ExecutorService 处理任务
     *
     * @param executor 线程池
     */
    public void start(ExecutorService executor) {
        this.init();

        // using count down latch
        this.latch = new AtomicInteger(this.size);

        // start working thread
        int n = this.size;
        while (n-- > 0) {
            executor.execute(this.runnable());
        }

        // if using provider
        if (this.provider != null) {
            this.provide();
        }
    }

    /**
     * 使用简单的线程工厂处理任务
     */
    public void start() {
        this.start(WorkTask.class.getName());
    }

    /**
     * 使用简单的线程工厂处理任务
     *
     * @param prefix 线程名称前缀
     */
    public void start(String prefix) {
        this.init();

        // using count down latch
        this.latch = new AtomicInteger(this.size);

        // executor
        CustomThreadFactory factory = new CustomThreadFactory(prefix);

        // start working thread
        int n = this.size;
        while (n-- > 0) {
            factory.newThread(this.runnable()).start();
        }

        // if using provider
        if (this.provider != null) {
            this.provide();
        }
    }

    /**
     * 增加元素，如果达到队列容量，则阻塞线程
     *
     * @param t 元素
     */
    public void put(U t) {
        try {
            this.queue.put(t);
        } catch (InterruptedException e) {
            throw new HandlerException("put element failed: ", e);
        }
    }

    /**
     * 添加哑元元素，runnable 读取到该元素，就会结束线程
     */
    public void putDumb() {
        for (WorkRunnable<T, U> run : list) {
            try {
                run.putDumb();
            } catch (InterruptedException e) {
                throw new HandlerException("put dumb element failed: ", e);
            }
        }
    }

    /**
     * count down latch
     * <p>
     * 子线程完成任务之后，会将计数减一，在计数降为 0 的时候，调用完成任务完成回调
     */
    public synchronized void countDown() {
        int n = latch.addAndGet(-1);
        if (n == 0 && this.finalized != null) {
            this.finalized.accept(this);
        } else if (n < 0) {
            throw new IndexOutOfBoundsException(Thread.currentThread() + " count down out of bounds: " + n);
        }
    }

    @Override
    public void close() {
        this.putDumb();
    }
}